package com.nexus.ai.rag;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.CompressionQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.TranslationQueryTransformer;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

/**
 * 检索增强
 * @Date 2025/9/29 03:51
 * @Author luzhengning
 **/
@Service
public class NexusAiRagService {

    @Autowired
    private VectorStore vectorStore;

    @Autowired
    private ChatClient chatClient;

    @Autowired
    private ChatMemory chatMemory;

    public String rag(String sessionId,String msg){
        //文字模板，选填
        String template=
                "以下为相关背景(上下文)信息，" +
                        "---------------------------" +
                        "{question_answer_context}" +
                        "---------------------------" +
                        "要求：回答需使用精准而简洁的回答，" +
                        "如果无法回答，请说明‘信息不足，暂时无法提供答案’，" +
                        "避免使用诸如‘根据上下文...’或‘提供的信息’这样的描述。" +
                        "问题：{query}";
        PromptTemplate customTemplate = PromptTemplate.builder().template(template).build();
        //元数据查询
        Filter.Expression filter=new Filter.Expression(
                Filter.ExpressionType.LT,//大于
                new Filter.Key("created_at"),
                new Filter.Value(System.currentTimeMillis())
        );
        //创建Advisor
        Advisor advisor= RetrievalAugmentationAdvisor.builder()
                //基本查询构建
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .vectorStore(vectorStore)
                        .similarityThreshold(0.1)   //相似度阈值
                        .topK(1)    //找出相似度最好的一条
                        //.filterExpression(filter)   //元数据查询
                        .build()
                )
                //构建查询压缩器
                .queryTransformers(
                        //解决查询格式不规范，术语含糊，词汇复杂或语言不支持等问题，使其更加精炼和明确
                        RewriteQueryTransformer.builder()
                        //创建聊天客户端的可变副本，保留原有配置的同时允许修改
                        .chatClientBuilder(this.chatClient.mutate())
                        .build(),
                        //构建查询翻译
                        TranslationQueryTransformer.builder()
                                .chatClientBuilder(chatClient.mutate())
                                .targetLanguage("Chinese")
                                .build(),
                        //构建上下文压缩,将对话历史记录压缩成一个独立的查询，能够准确反映对话的核心要点
                        CompressionQueryTransformer.builder()
                                .chatClientBuilder(chatClient.mutate())
                                .build()
                )
                //构建多角度提问，将一个问题转化为多个语义各异的变体，以涵盖不同的角度，对于获取更多背景以及提高查询精准度非常有用
                .queryExpander(MultiQueryExpander.builder()
                        .chatClientBuilder(chatClient.mutate())
                        .numberOfQueries(3)     //生成问题的个数
                        .includeOriginal(true)  //是否包含原始问题
                        .build()
                )
                //允许上下文为空，如果检索到上下文为空(Rag库中不存在该类问题的答案)，它指示模型不回答用户
                .queryAugmenter(ContextualQueryAugmenter.builder()
                        .allowEmptyContext(true)    //允许上下文为空,不包含模型数据的内容
                        .build()
                )
                //文档预处理/结果重排序，是指对从向量数据库检索到的文档安相关性重新排序和过滤
                .documentPostProcessors((query, documents) -> {
                    documents.forEach(doc->{
                        //执行业务逻辑，修改文档内容
                        doc.getMetadata().put("score",0.5);
                    });
                    //结果重新排序，并返回处理后的文档
                    return documents;
                })
                .build();

        //对话记忆
        MessageChatMemoryAdvisor chatMemroyAdvisor = MessageChatMemoryAdvisor.builder(chatMemory).conversationId(sessionId).build();

        ChatClient.CallResponseSpec call = chatClient.prompt()
                .user(msg)
                //使用Advisor实现Rag增强查询
                .advisors(
                        //查询
                        advisor
                        //对话记忆
                        ,chatMemroyAdvisor
                )
                .call();
        String content = call.content();
        return content;
    }
}
